from os import listdir
from collections import Counter
from itertools import chain
from numpy import array
from sklearn.externals import joblib
from sklearn.naive_bayes import MultinomialNB
from get_words_from_file import getWordsFromFile

# 存放所有文件中的单词
# 每个元素是一个子列表，其中存放一个文件中的所有单词
allWords = []
def getTopNWords(topN):
    # 按文件编号顺序处理当前文件夹中所有记事本文件
    # 训练集中共151封邮件内容，0.txt到126.txt是垃圾邮件内容
    # 127.txt到150.txt为正常邮件内容
    txtFiles = [str(i)+'.txt' for i in range(151)]
    # 获取训练集中所有邮件中的全部单词
    for txtFile in txtFiles:
        allWords.append(getWordsFromFile(txtFile))
    # 获取并返回出现次数最多的前topN个单词
    freq = Counter(chain(*allWords))
    return [w[0] for w in freq.most_common(topN)]

# 全部训练集中出现次数最多的前600个单词
topWords = getTopNWords(600)

# 获取特征向量，前600个单词的每个单词在每个邮件中出现的频率
vectors = []
for words in allWords:
    temp = list(map(lambda x: words.count(x), topWords))
    vectors.append(temp)
vectors = array(vectors)
# 训练集中每个邮件的标签，1表示垃圾邮件，0表示正常邮件
labels = array([1]*127 + [0]*24)

# 创建模型，使用已知训练集进行训练
model = MultinomialNB()
model.fit(vectors, labels)

joblib.dump(model, "垃圾邮件分类器.pkl")
print('保存模型和训练结果成功。')
with open('topWords.txt', 'w', encoding='utf8') as fp:
    fp.write(','.join(topWords))
print('保存topWords成功。')
